import os  
import shutil
from tqdm import tqdm
import numpy as np
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
from pytorch_lightning import seed_everything
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import datasets, models, transforms
from torchvision.utils import save_image
from torch.utils.data import DataLoader
from utils import *

from torch.utils.tensorboard import SummaryWriter

%load_ext autoreload
%autoreload 1

configs = {   
    "model_name" : "CVAE",
    "exp" : "1",  
    "latent_dim" : 64,
    "batch_size" : 64,
    "num_epochs" : 50,
    "lr" : 1e-3,
    "scheduler" : "ReduceLROnPlateau",
    "use_scheduler" : True,
    "lambda_kld" : 1e-3,
    }

dataset_root = './data/AFHQ/'

transform = transforms.Compose([transforms.Resize((64,64)),
                                      transforms.ToTensor(),
                                      transforms.Normalize([0.5]*3 , [0.5]*3)])

BS = configs["batch_size"]

train_dataset = datasets.ImageFolder(root= dataset_root+'train', transform= transform )
test_dataset = datasets.ImageFolder(root= dataset_root+'test', transform= transform )

# print(train_dataset.classes)  
print(train_dataset.class_to_idx)  

train_loader = DataLoader(dataset= train_dataset, 
                          batch_size= BS, 
                          shuffle= True, 
                          drop_last= True )

test_loader = DataLoader(dataset= test_dataset, 
                          batch_size= BS, 
                          shuffle= False, 
                          drop_last= True )

{'cat': 0, 'dog': 1, 'wild': 2}

imgs, labels = next(iter(train_loader))
fig, ax = plt.subplots(1, 8)
fig.set_size_inches(18, 5)
for i in range(8):
    img_np = imgs[i].permute(1, 2, 0).numpy()
    # Denormalize the image
    img_np = (img_np * 0.5 + 0.5).clip(0, 1)
    ax[i].imshow(img_np)
    ax[i].axis("off")
    ax[i].set_title(f"{train_dataset.classes[labels[i]]}")
plt.tight_layout()
plt.show()

imgs.shape

torch.Size([64, 3, 64, 64])

class CVAE(nn.Module):
    def __init__(self, in_channels=3, latent_dim=64):
        super(CVAE, self).__init__()
        
        self.latent_dim = latent_dim
        self.in_channels = in_channels
        
        # Encoder
        self.encoder = self.make_encoder()
        
        # Fully connected layers for mu and sigma
        self.fc_mu = nn.Linear(2048, latent_dim)
        self.fc_sigma = nn.Linear(2048, latent_dim)
        
        '''
        Also called projection layer, projects latent vector to  original dimensionality.
        maps the latent variable into a space compatible with the decoder input
        '''
        self.decoder_input = nn.Sequential(
            nn.Linear(latent_dim, 2048),
            nn.LeakyReLU(0.2)
        )
        
        # Decoder
        self.decoder = self.make_decoder()

    def make_encoder(self):
        
        return nn.Sequential(
            # Input: (3, 64, 64)
            nn.Conv2d(self.in_channels, 16, kernel_size=4, stride=2, padding=1),  # (16, 32, 32)
            nn.BatchNorm2d(16),
            nn.LeakyReLU(0.2),
            
            nn.Conv2d(16, 32, kernel_size=4, stride=2, padding=1),  # (32, 16, 16)
            nn.BatchNorm2d(32),
            nn.LeakyReLU(0.2),
            
            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),  # (64, 8, 8)
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2),
            
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),  # (128, 4, 4)
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2),
            
            nn.Flatten()  # 128 * 4 * 4 = 2048
        )
        
    def make_decoder(self):
        return nn.Sequential(
            # Input: (128, 4, 4)
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),  # (64, 8, 8)
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2),
            
            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),  # (32, 16, 16)
            nn.BatchNorm2d(32),
            nn.LeakyReLU(0.2),
            
            nn.ConvTranspose2d(32, 16, kernel_size=4, stride=2, padding=1),  # (16, 32, 32)
            nn.BatchNorm2d(16),
            nn.LeakyReLU(0.2),
            
            nn.ConvTranspose2d(16, self.in_channels, kernel_size=4, stride=2, padding=1),  # (3, 64, 64)
            nn.Sigmoid()
        )

    def reparameterize(self, mu, log_var):
        """Reparameterization trick"""
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        return mu + eps * std
        
    def forward(self, x):
        # Encode
        x_encoded = self.encoder(x)
        
        # Get mu and log_var
        mu = self.fc_mu(x_encoded)  # [batch_size, latent_dim]
        log_var = self.fc_sigma(x_encoded)  # [batch_size, latent_dim]
        
        # Reparameterization trick
        z = self.reparameterize(mu, log_var)  # [batch_size, latent_dim]
        
        # Project and reshape for decoder
        z = self.decoder_input(z)
        z = z.view(-1, 128, 4, 4)
        x_hat = self.decoder(z)
        
        return x_hat, (z, mu, log_var)

device = "cuda" if torch.cuda.is_available() else "cpu"

seed_everything(42)

model = CVAE(latent_dim=configs["latent_dim"]).to(device)

Seed set to 42

model

CVAE(
  (encoder): Sequential(
    (0): Conv2d(3, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): LeakyReLU(negative_slope=0.2)
    (3): Conv2d(16, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (4): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (5): LeakyReLU(negative_slope=0.2)
    (6): Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (7): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (8): LeakyReLU(negative_slope=0.2)
    (9): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (10): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (11): LeakyReLU(negative_slope=0.2)
    (12): Flatten(start_dim=1, end_dim=-1)
  )
  (fc_mu): Linear(in_features=2048, out_features=64, bias=True)
  (fc_sigma): Linear(in_features=2048, out_features=64, bias=True)
  (decoder_input): Sequential(
    (0): Linear(in_features=64, out_features=2048, bias=True)
    (1): LeakyReLU(negative_slope=0.2)
  )
  (decoder): Sequential(
    (0): ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): LeakyReLU(negative_slope=0.2)
    (3): ConvTranspose2d(64, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (4): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (5): LeakyReLU(negative_slope=0.2)
    (6): ConvTranspose2d(32, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (7): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (8): LeakyReLU(negative_slope=0.2)
    (9): ConvTranspose2d(16, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
    (10): Sigmoid()
  )
)

count_model_params(model.encoder)

173520

count_model_params(model.decoder)

173139

count_model_params(model.fc_mu) + count_model_params(model.fc_sigma)

262272

count_model_params(model)

742051

optimizer = torch.optim.AdamW(model.parameters(), lr=configs["lr"], weight_decay=1e-5)

scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=7, factor=0.5, verbose=True)

/home/user/soltania1/.local/lib/python3.8/site-packages/torch/optim/lr_scheduler.py:60: UserWarning: The verbose parameter is deprecated. Please use get_last_lr() to access the learning rate.
  warnings.warn(

def makedires(configs):
    model_name = configs["model_name"]+configs["exp"]+f"_KLD_{configs['lambda_kld']}"
    savepath = f"imgs/{model_name}"
    if os.path.exists(savepath):
        shutil.rmtree(savepath)
    os.makedirs(savepath,exist_ok=True)

    TBOARD_LOGS = os.path.join(os.getcwd(), "tboard_logs", model_name+f"_KLD_{configs['lambda_kld']}")
    if os.path.exists(TBOARD_LOGS):
        shutil.rmtree(TBOARD_LOGS)
    os.makedirs(TBOARD_LOGS)
    writer = SummaryWriter(TBOARD_LOGS)
    return savepath, writer

savepath, writer = makedires(configs)

def vae_loss_function(recons, target, mu, log_var, lambda_kld=configs["lambda_kld"]):
    """
    Combined loss function for joint optimization of 
    reconstruction and ELBO
    """
    recons_loss = F.mse_loss(recons, target)
    kld = (-0.5 * (1 + log_var - mu**2 - log_var.exp()).sum(dim=1)).mean(dim=0)  # closed-form solution of KLD in Gaussian
    loss = recons_loss + lambda_kld * kld

    return loss, (recons_loss, kld)

model_name = configs["model_name"]+configs["exp"]+f"_KLD_{configs['lambda_kld']}"

train_loss, val_loss, loss_iters, val_loss_recons, val_loss_kld = train_model(
        model=model, 
        model_name=model_name,
        optimizer=optimizer,
        scheduler=scheduler if configs["use_scheduler"] else None, 
        criterion=vae_loss_function,
        lambda_kld=configs["lambda_kld"],
        train_loader=train_loader, 
        valid_loader=test_loader, 
        num_epochs=configs["num_epochs"],
        save_frequency= 10, 
        savepath=savepath,
        writer=writer
    )

Epoch 1 Iter 224: loss 0.22007. : 100%|██████████| 224/224 [01:43<00:00,  2.17it/s]

    Train loss: 0.29521
    Valid loss: 0.60873
       Valid loss recons: 0.60871
       Valid loss KL-D:   0.01345

Epoch 2 Iter 224: loss 0.19876. : 100%|██████████| 224/224 [01:43<00:00,  2.17it/s]
Epoch 3 Iter 224: loss 0.20027. : 100%|██████████| 224/224 [01:41<00:00,  2.21it/s]

    Train loss: 0.20983
    Valid loss: 0.21243
       Valid loss recons: 0.20206
       Valid loss KL-D:   10.3759

Epoch 4 Iter 224: loss 0.21310. : 100%|██████████| 224/224 [01:43<00:00,  2.17it/s]
Epoch 5 Iter 224: loss 0.20445. : 100%|██████████| 224/224 [01:39<00:00,  2.24it/s]

    Train loss: 0.20653
    Valid loss: 0.20849
       Valid loss recons: 0.19666
       Valid loss KL-D:   11.8341

Epoch 6 Iter 224: loss 0.20707. : 100%|██████████| 224/224 [01:40<00:00,  2.24it/s]
Epoch 7 Iter 224: loss 0.18920. : 100%|██████████| 224/224 [01:41<00:00,  2.21it/s]

    Train loss: 0.20498
    Valid loss: 0.20621
       Valid loss recons: 0.19423
       Valid loss KL-D:   11.98574

Epoch 8 Iter 224: loss 0.21119. : 100%|██████████| 224/224 [01:40<00:00,  2.24it/s]
Epoch 9 Iter 224: loss 0.19564. : 100%|██████████| 224/224 [01:40<00:00,  2.22it/s]

    Train loss: 0.20419
    Valid loss: 0.20548
       Valid loss recons: 0.19329
       Valid loss KL-D:   12.1849

def save_config(configs):
    model_name = configs["model_name"]+configs["exp"]+f"_KLD_{configs['lambda_kld']}"
    configs_dir = f"./configs/{model_name}/"
    if not os.path.exists(configs_dir):
        os.makedirs(configs_dir,exist_ok=True)
    configs_path = configs_dir + "/config.yaml"

    with open(configs_path, 'w') as f:
        yaml.dump(configs, f)

epoch = configs["num_epochs"]
lambda_kld = configs["lambda_kld"]

save_model(model, model_name, optimizer, epoch = epoch,lambda_kld = lambda_kld, stats = configs )
save_config(configs)

img_vs_recons(model, test_loader, device)

def compute_encoder_output_size(batch_size, input_shape, model):
    """
    Compute the un/flattened output size of the encoder given an input shape.
    """
    with torch.no_grad():

        enc_input = torch.zeros(batch_size, *input_shape).to(device)  # (BS, 3, 64, 64)
        
        # Get encoder output
        output = model.encoder(enc_input)
        
        enc_out_shape = output.view(batch_size,-1,4,4).shape[1:]
        # Get flattened size
        flattened_size = output.view(1, -1).shape[1]
        return enc_out_shape, flattened_size

compute_encoder_output_size(BS, (3, 64, 64), model)

(torch.Size([128, 4, 4]), 131072)

def inference(configs, model):
    model_name = configs["model_name"]+configs["exp"]+f"_KLD_{configs['lambda_kld']}"
    BS = configs["batch_size"]
    if not os.path.exists(f"imgs/inference/{model_name}"):
        os.makedirs(f"imgs/inference/{model_name}")

    latent_dim = configs["latent_dim"]

    enc_output_shape, _ = compute_encoder_output_size(BS, (3, 64, 64), model)
    # print(enc_output_shape)
    with torch.no_grad():
        for i in range(5):
            z = torch.randn(BS, latent_dim).to(device)

            z = model.decoder_input(z)
            z = z.view(-1, *enc_output_shape)

            recons = model.decoder(z)
            recons = recons.view(BS, 3, 64, 64)
            save_image(recons, f"imgs/inference/{model_name}/inference_{i}.png")
            return recons

def plot_recons(recons):
    plt.figure(figsize=(8*2, 4*2))
    for i in range(32):
        plt.subplot(4,8,i+1)
        recon_img = denormalize_images(recons[i].cpu())
        # recon_img = recons[i].cpu()
        recon_img = recon_img.permute(1, 2, 0)
        plt.imshow(recon_img)  
        plt.axis("off")
    plt.tight_layout()
    plt.show()
    
    
plot_recons(inference(configs,model))

def vis_latent(test_loader, model, test_dataset):
    imgs_flat, latents, labels = [], [], []
    with torch.no_grad():
        for imgs, lbls in test_loader:
            imgs = imgs.to(device)
            _, (z, _, _) = model(imgs)
            imgs_flat.append(imgs.cpu().view(imgs.shape[0],-1))
            latents.append(z.cpu())
            labels.append(lbls)
            
    imgs_flat = np.concatenate(imgs_flat)    
    latents = np.concatenate(latents)
    labels = np.concatenate(labels)

    latents_reshaped = latents.reshape(latents.shape[0], -1)

    pca_imgs = PCA(n_components=2).fit_transform(imgs_flat)
    pca_latents = PCA(n_components=2).fit_transform(latents_reshaped)

    N = 2000
    fig,ax = plt.subplots(1,2,figsize=(26,8))
    display_projections(pca_imgs[:N], labels[:N], ax=ax[0], legend=test_dataset.classes)
    ax[0].set_title("PCA Proj. of Images")
    display_projections(pca_latents[:N], labels[:N], ax=ax[1], legend=test_dataset.classes)
    ax[1].set_title("Encoded Representations")
    plt.show()

vis_latent(test_loader, model, test_dataset)

@torch.no_grad()
def plot_reconstructed(model, xrange=(-3, 3), yrange=(-10, 10), N=12):
    """
    Sampling equispaced points from the latent space given the xrange and yrange, 
    decoding latents and visualizing distribution of the space
    """
    # Project points to decoder input dimension (same as in forward pass)
    SIZE = 64  # Image size
    grid = np.empty((N*SIZE, N*SIZE, 3))  # 3 channels for RGB
    
    for i, y in enumerate(np.linspace(*yrange, N)):
        for j, x in enumerate(np.linspace(*xrange, N)):
            # mean
            mu = torch.zeros(model.latent_dim, device=device)
            mu[0] = x
            mu[1] = y
            
            # standard deviation
            # Create a vector of stds the same shape as mu
            sigma = 1 
            std = torch.full_like(mu, sigma)
            z = torch.normal(mean=mu, std=std)
            
            # Passing through the decoder 
            z = model.decoder_input(z)
            z = z.view(-1, 128, 4, 4)  # Reshape to match decoder input
            
            # Getting recons
            x_hat = model.decoder(z)
            
            # To visualize
            x_hat = x_hat.squeeze(0).cpu()  # Remove batch dimension
            x_hat = x_hat.permute(1, 2, 0)  # (C,H,W) to (H,W,C)
            recon_img = x_hat.numpy()
            
            # Prepare the image for visualization
            recon_img = denormalize_images(recon_img[i])
            
            grid[(N-1-i)*SIZE:(N-i)*SIZE, j*SIZE:(j+1)*SIZE] = recon_img
           
    plt.figure(figsize=(12,20))
    plt.imshow(grid, extent=[*yrange, *xrange])
    plt.axis("off")
    plt.show()

plot_reconstructed(model, xrange=(-10, 10), yrange=(-10, 10), N=15)

seed_everything(42)

model = CVAE(latent_dim=configs["latent_dim"]).to(device)

optimizer = torch.optim.AdamW(model.parameters(), lr=configs["lr"], weight_decay=1e-5)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=7, factor=0.5, verbose=True)
model_name = configs["model_name"]+configs["exp"]+f"_KLD_{configs['lambda_kld']}"
savepath, writer = makedires(configs)

train_loss, val_loss, loss_iters, val_loss_recons, val_loss_kld = train_model(
        model=model, 
        model_name=model_name,
        optimizer=optimizer,
        scheduler=scheduler if configs["use_scheduler"] else None, 
        criterion=vae_loss_function,
        lambda_kld=configs["lambda_kld"],
        train_loader=train_loader, 
        valid_loader=test_loader, 
        num_epochs=configs["num_epochs"],
        save_frequency= 10, 
        savepath=savepath,
        writer=writer
    )

epoch = configs["num_epochs"]
lambda_kld = configs["lambda_kld"]

save_model(model, model_name, optimizer, epoch = epoch,lambda_kld = lambda_kld, stats = configs )
save_config(configs)

Seed set to 42

/home/user/soltania1/.local/lib/python3.8/site-packages/torch/optim/lr_scheduler.py:60: UserWarning: The verbose parameter is deprecated. Please use get_last_lr() to access the learning rate.
  warnings.warn(

  0%|                                                                                        | 0/224 [00:00<?, ?it/s]

Epoch 1 Iter 1: loss 0.53192. :   0%|                                                        | 0/224 [00:00<?, ?it/s]

Epoch 1 Iter 1: loss 0.53192. :   0%|▏                                               | 1/224 [00:00<02:01,  1.84it/s]

Epoch 1 Iter 2: loss 0.52736. :   0%|▏                                               | 1/224 [00:01<02:01,  1.84it/s]

Epoch 1 Iter 2: loss 0.52736. :   1%|▍                                               | 2/224 [00:01<01:51,  1.98it/s]

Epoch 1 Iter 3: loss 0.56549. :   1%|▍                                               | 2/224 [00:01<01:51,  1.98it/s]

Epoch 1 Iter 3: loss 0.56549. :   1%|▋                                               | 3/224 [00:01<01:46,  2.07it/s]

Epoch 1 Iter 4: loss 0.49826. :   1%|▋                                               | 3/224 [00:01<01:46,  2.07it/s]

img_vs_recons(model, test_loader, device)

plot_recons(inference(configs,model))

vis_latent(test_loader, model, test_dataset)

plot_reconstructed(model, xrange=(-2000, 2000), yrange=(-2000, 2000), N=15)

seed_everything(42)

model = CVAE(latent_dim=configs["latent_dim"]).to(device)

optimizer = torch.optim.AdamW(model.parameters(), lr=configs["lr"], weight_decay=1e-5)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=7, factor=0.5, verbose=True)
model_name = configs["model_name"]+configs["exp"]+f"_KLD_{configs['lambda_kld']}"
savepath, writer = makedires(configs)

train_loss, val_loss, loss_iters, val_loss_recons, val_loss_kld = train_model(
        model=model, 
        model_name=model_name,
        optimizer=optimizer,
        scheduler=scheduler if configs["use_scheduler"] else None, 
        criterion=vae_loss_function,
        lambda_kld=configs["lambda_kld"],
        train_loader=train_loader, 
        valid_loader=test_loader, 
        num_epochs=configs["num_epochs"],
        save_frequency= 10, 
        savepath=savepath,
        writer=writer
    )

epoch = configs["num_epochs"]
lambda_kld = configs["lambda_kld"]

save_model(model, model_name, optimizer, epoch = epoch,lambda_kld = lambda_kld, stats = configs )
save_config(configs)

Seed set to 42

/home/user/soltania1/.local/lib/python3.8/site-packages/torch/optim/lr_scheduler.py:60: UserWarning: The verbose parameter is deprecated. Please use get_last_lr() to access the learning rate.
  warnings.warn(

  0%|                                                                                        | 0/224 [00:00<?, ?it/s]

Epoch 1 Iter 1: loss 0.54591. :   0%|                                                        | 0/224 [00:00<?, ?it/s]

Epoch 1 Iter 1: loss 0.54591. :   0%|▏                                               | 1/224 [00:00<02:09,  1.72it/s]

Epoch 1 Iter 2: loss 1.34620. :   0%|▏                                               | 1/224 [00:01<02:09,  1.72it/s]

Epoch 1 Iter 2: loss 1.34620. :   1%|▍                                               | 2/224 [00:01<01:57,  1.89it/s]

Epoch 1 Iter 3: loss 0.64037. :   1%|▍                                               | 2/224 [00:01<01:57,  1.89it/s]

Epoch 1 Iter 3: loss 0.64037. :   1%|▋                                               | 3/224 [00:01<01:52,  1.96it/s]

Epoch 1 Iter 4: loss 0.57654. :   1%|▋                                               | 3/224 [00:02<01:52,  1.96it/s]

img_vs_recons(model, test_loader, device)

plot_recons(inference(configs,model))

vis_latent(test_loader, model, test_dataset)

plot_reconstructed(model, xrange=(-400, 400), yrange=(-400, 400), N=15)

seed_everything(42)

model = CVAE(latent_dim=configs["latent_dim"]).to(device)

optimizer = torch.optim.AdamW(model.parameters(), lr=configs["lr"], weight_decay=1e-5)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=7, factor=0.5, verbose=True)
model_name = configs["model_name"]+configs["exp"]+f"_KLD_{configs['lambda_kld']}"
savepath, writer = makedires(configs)

train_loss, val_loss, loss_iters, val_loss_recons, val_loss_kld = train_model(
        model=model, 
        model_name=model_name,
        optimizer=optimizer,
        scheduler=scheduler if configs["use_scheduler"] else None, 
        criterion=vae_loss_function,
        lambda_kld=configs["lambda_kld"],
        train_loader=train_loader, 
        valid_loader=test_loader, 
        num_epochs=configs["num_epochs"],
        save_frequency= 10, 
        savepath=savepath,
        writer=writer
    )

epoch = configs["num_epochs"]
lambda_kld = configs["lambda_kld"]

save_model(model, model_name, optimizer, epoch = epoch,lambda_kld = lambda_kld, stats = configs )
save_config(configs)

Seed set to 42

/home/user/soltania1/.local/lib/python3.8/site-packages/torch/optim/lr_scheduler.py:60: UserWarning: The verbose parameter is deprecated. Please use get_last_lr() to access the learning rate.
  warnings.warn(

  0%|                                                                                        | 0/224 [00:00<?, ?it/s]

Epoch 1 Iter 1: loss 0.53206. :   0%|                                                        | 0/224 [00:00<?, ?it/s]

Epoch 1 Iter 1: loss 0.53206. :   0%|▏                                               | 1/224 [00:00<02:10,  1.71it/s]

Epoch 1 Iter 2: loss 0.53072. :   0%|▏                                               | 1/224 [00:01<02:10,  1.71it/s]

Epoch 1 Iter 2: loss 0.53072. :   1%|▍                                               | 2/224 [00:01<01:59,  1.86it/s]

Epoch 1 Iter 3: loss 0.55081. :   1%|▍                                               | 2/224 [00:01<01:59,  1.86it/s]

Epoch 1 Iter 3: loss 0.55081. :   1%|▋                                               | 3/224 [00:01<01:54,  1.93it/s]

Epoch 1 Iter 4: loss 0.51249. :   1%|▋                                               | 3/224 [00:02<01:54,  1.93it/s]

img_vs_recons(model, test_loader, device)

plot_recons(inference(configs,model))

vis_latent(test_loader, model, test_dataset)

plot_reconstructed(model, xrange=(-50, 50), yrange=(-50, 50), N=15)

λ (KLD weight)	Effect
`0.0`	Pure autoencoder (no latent regularization)
`1e-4`	Very weak KLD → good reconstructions, weak latent structure
`1e-3`	Balanced trade-off
`1e-2`	Strong KLD → worse recons, better sampling and latent space

Assignment 4¶

Team name:¶

Members:¶

Experiment with `KLD = 1e_3`¶

$Note\ on\ network\ architecture$¶

Saving model and training configs¶

Inference¶

Latent space visualization¶

Interpolation¶

Investigating the importance of the KL-divergence¶

Experiment with `KLD = 0.0`¶

Images vs recons¶

Inference¶

Latent space visualization¶

Interpolation¶

Experiment with `KLD = 1e_2`¶

Images vs recons¶

Inference¶

Latent space visualization¶

Interpolation¶

Experiment with `KLD = 1e_4`¶

Images vs recons¶

Inference¶

Latent space visualization¶

Interpolation¶

$Analysis$¶

Extra Point:¶

Assignment 4¶

Team name:¶

Members:¶

Experiment with KLD = 1e_3¶

$Note\ on\ network\ architecture$¶

Saving model and training configs¶

Inference¶

Latent space visualization¶

Interpolation¶

Investigating the importance of the KL-divergence¶

Experiment with KLD = 0.0¶

Images vs recons¶

Inference¶

Latent space visualization¶

Interpolation¶

Experiment with KLD = 1e_2¶

Images vs recons¶

Inference¶

Latent space visualization¶

Interpolation¶

Experiment with KLD = 1e_4¶

Images vs recons¶

Inference¶

Latent space visualization¶

Interpolation¶

$Analysis$¶

Extra Point:¶

Experiment with `KLD = 1e_3`¶

Experiment with `KLD = 0.0`¶

Experiment with `KLD = 1e_2`¶

Experiment with `KLD = 1e_4`¶